"""Cosine classifier for continual learning tasks."""

import torch
import torch.nn as nn
import torch.nn.functional as F


class CosineClassifier(nn.Module):
    """Basic cosine similarity classifier.

    Computes classification logits using cosine similarity between input features
    and learnable class prototypes. This approach provides better feature separation
    and is particularly effective for continual learning scenarios.

    Args:
        input_dim: Dimension of input features
        num_classes: Number of output classes
        temperature: Temperature scaling factor for logits (default: 1.0)
        normalize_features: Whether to L2-normalize input features (default: True)

    Example:
        >>> classifier = CosineClassifier(input_dim=512, num_classes=10)
        >>> features = torch.randn(32, 512)
        >>> logits = classifier(features)
        >>> print(logits.shape)  # torch.Size([32, 10])
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        temperature: float = 0.1,
    ) -> None:
        super().__init__()

        self.in_features = in_features
        self.out_features = out_features

        # Learnable class prototypes
        self.weight = nn.Parameter(
            torch.empty(out_features, in_features, dtype=torch.float32)
        )

        # Temperature parameter (fixed)
        self.register_buffer("temperature", torch.tensor(temperature))

        self._initialize_parameters()

    def _initialize_parameters(self) -> None:
        """Initialize weight parameters using Xavier uniform initialization."""
        nn.init.xavier_uniform_(self.weight)

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """Forward pass computing cosine similarity-based logits.

        Args:
            features: Input features of shape (batch_size, input_dim)

        Returns:
            Classification logits of shape (batch_size, num_classes)
        """
        features = F.normalize(features, p=2, dim=1)
        normalized_weight = F.normalize(self.weight, p=2, dim=1)
        cosine_similarity = F.linear(features, normalized_weight)

        # Apply temperature scaling
        logits = cosine_similarity / self.temperature

        return logits

    def extra_repr(self) -> str:
        """Return string representation for debugging."""
        return (
            f"in_features={self.in_features}, "
            f"out_features={self.out_features}, "
            f"temperature={self.temperature.item():.3f}"
        )
